import pandas as pd
import os,glob,math
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.lines as mlines
import matplotlib.dates as mdates
import cartopy.crs as ccrs
from cartopy.mpl.ticker import LongitudeFormatter,LatitudeFormatter
import cartopy.feature as cfeature
import shapely.geometry as sgeom
from datetime import datetime,timedelta
from tqdm import tqdm
from global_land_mask import globe
from gsj_typhoon import tydat,see,count_rapidgrow,tydat_CMA,average_datetime,split_str_id,load_land_polygons,detect_landfall
from geopy.distance import geodesic
import matplotlib.ticker as ticker
import seaborn as sns
from typlot.config.global_config import *


all_ini_time_mode = ['12','00','00_12']
names = ['mojie_28','dusurui_16','gaemi_09','haikui_38','kangni_54','shantuo_44','saola_25','koinu_49']
#names = ['dusurui_16']
tynames,tyids = split_str_id(names)
draw_obs_opt = True
all_obs_baseline= ['land','RI']  # ‘land’ 'RI'
draw_opt = False


for ini_time_mode in all_ini_time_mode:
    for obs_baseline in all_obs_baseline:
          
        for ty,tyid in zip(tynames,tyids):
            ''' init '''
            RIstd = 7
            dt = timedelta(days=1)
            directory = os.path.join(global_ensdir,f"{ty}_{tyid}")
            dates_name = sorted(os.listdir(directory))
            dates_name = [i for i in dates_name if i[-2:] in ini_time_mode ]
            dates_paths = [os.path.join(directory,date) for date in dates_name if date[-2:] in ini_time_mode]
            obs_path = os.path.join(global_obsdir,f"{ty}_CMAobs.txt")
            pic_savepath=os.path.join(global_picdir,f'{obs_baseline}_heatmap',ini_time_mode,'start_day_distance')
            os.makedirs(pic_savepath,exist_ok=True)
            
            
            ''' calc 2d numbers '''
            # 遍历所有起报时间
            max_delta_time = timedelta(days=0)
            for j,i in enumerate(dates_paths):
                mem_paths = glob.glob( os.path.join(i,"TRACK_ID_*") )
                # sorted_mem_paths = sorted(mem_paths, key=lambda x: int(x.split('_')[-1]))
                start_time = datetime.strptime(dates_name[j],"%Y%m%d%H")
                # 遍历所有成员
                for single_mem_path in mem_paths:
                    single_mem = tydat(single_mem_path,RIstd) 
                    delta_time = single_mem.time[-1] - start_time
                    if delta_time>max_delta_time:
                        max_delta_time = delta_time
            max_dt_num = math.ceil(max_delta_time/dt) 
            
            
            ''' 求各起报时间的图 '''
            tyobs = tydat_CMA(obs_path)
            rcd_obs = count_rapidgrow(RIstd, tyobs.umax, tyobs.time)
            num_result = np.zeros(( max_dt_num,len(dates_name) ))
            dis_result = np.zeros(( max_dt_num,len(dates_name) ))
            dis_var_result = np.zeros(( max_dt_num,len(dates_name) ))
            
            
            for i in tqdm( range(len(dates_paths)), desc=f'processing  {ty} {obs_baseline}_{ini_time_mode}'):
                '''读取这个起报时间的所有集合'''
                mem_paths = glob.glob( os.path.join(dates_paths[i],"TRACK_ID_*") )
                sorted_mem_paths = sorted(mem_paths, key=lambda x: int(x.split('_')[-1]))
                ty_mem = [tydat(i,RIstd) for i in sorted_mem_paths]
                start_dt = datetime.strptime( dates_name[i],"%Y%m%d%H" )
                for j in range(max_dt_num):
                    t0 = j
                    t1 = j+1
                    t0_dt = start_dt + timedelta(days=t0)
                    t1_dt = start_dt + timedelta(days=t1)
                    '''如果观测在这两天没有RI，就直接跳过'''
                    idx_obs_2d = (tyobs.time>=t0_dt)&(tyobs.time<t1_dt)
                    idx_obs_2d_RI = (rcd_obs[idx_obs_2d]==1)
                    if np.sum(rcd_obs[idx_obs_2d]) ==0.:
                        ''' only record numbers here '''
                        for k in range(len(ty_mem)):
                            '''cycle through all members : 如果这个成员在这两天没有RI,就跳过'''
                            idx_mem_2d    = (ty_mem[k].time>=t0_dt)&(ty_mem[k].time<t1_dt)
                            idx_mem_2d_RI = (ty_mem[k].num_rapidgrow()[idx_mem_2d]==1)
                            num_result[j,i] += np.sum(idx_mem_2d_RI)
                        continue
                    else:
                        '''如果2d内有多个RI,取平均'''
                        time_obs_ave = average_datetime(tyobs.time[idx_obs_2d][idx_obs_2d_RI])
                        lon_obs_ave  = np.mean(tyobs.lon[idx_obs_2d][idx_obs_2d_RI])
                        lat_obs_ave  = np.mean(tyobs.lat[idx_obs_2d][idx_obs_2d_RI])
                        '''开始计算'''
                        num = 0 
                        dis_list = [] 
        
                        for k in range(len(ty_mem)):
                            '''cycle through all members : 如果这个成员在这两天没有RI,就跳过'''
                            idx_mem_2d    = (ty_mem[k].time>=t0_dt)&(ty_mem[k].time<t1_dt)
                            idx_mem_2d_RI = (ty_mem[k].num_rapidgrow()[idx_mem_2d]==1)
                            time_mem = ty_mem[k].time[idx_mem_2d][idx_mem_2d_RI]
                            lat_mem  = ty_mem[k].lat[idx_mem_2d][idx_mem_2d_RI]
                            lon_mem  = ty_mem[k].lon[idx_mem_2d][idx_mem_2d_RI]
                            if np.sum(idx_mem_2d_RI) ==0.:
                                continue
                            else:
                                num+=np.sum(idx_mem_2d_RI)
                                for l in range(len(lon_mem)):
                                    dis_list.append(geodesic((lat_obs_ave,lon_obs_ave),(lat_mem[l],lon_mem[l])).kilometers)
        
                    if num==0:
                        continue
                    num_result[j,i] = num
                    dis_result[j,i] = np.array(dis_list).mean()
                    dis_var_result[j,i] = np.array(dis_list).std()
        
                    
            if draw_obs_opt == True : 
               ''' init obsCMA; RI points reord '''
               obs = tydat_CMA(obs_path) 
               draw_data_obs = np.zeros( (len(dates_name),max_dt_num),dtype=int )
               tt,dd = draw_data_obs.shape
               
               if obs_baseline == 'RI':
                    record_obs = count_rapidgrow(RIstd,obs.umax, obs.time)
                    for t in range(tt):
                       #预报时效
                       for d in range(dd):
                           try:
                               start_time = datetime.strptime(dates_name[t],"%Y%m%d%H")
                               cal_time = start_time + d*timedelta(days=1)
                               calend_time = cal_time + timedelta(days=1)
                               # 对齐
                               trange = (obs.time>=cal_time) & (obs.time<calend_time)
                               draw_data_obs[t][d] = 1 if np.sum(record_obs[trange])>0 else 0 
                           except Exception as e:
                               print('line 132',e)
               elif obs_baseline=='land':
                    # get  land  time
                    land_polys=load_land_polygons(os.path.join(global_shpdir,'China','bou1_4p.shp'),os.path.join(global_shpdir,'China','bou1_4p.dbf'))
                    flags,id1 = detect_landfall(obs.lat, obs.lon, land_polys)
                    land_time = obs.time[id1]
                    # identify  position
                    for t in range(tt):
                        for d in range(dd):
                            try:
                                start_time = datetime.strptime(dates_name[t],"%Y%m%d%H")
                                cal_time = start_time + d*timedelta(days=1)
                                calend_time = cal_time + timedelta(days=1)
                                trange = (obs.time>=cal_time) & (obs.time<calend_time)
                                # 填充
                                draw_data_obs[t][d] = 1 if land_time in obs.time[trange] else 0
                            except Exception as e:
                                print('line 149',e)
                    
        
            
        
            var = dis_result
            annot_str = np.empty_like(var[:15], dtype=object)
            for i in range(annot_str.shape[0]):
                for j in range(annot_str.shape[1]):
                    annot_str[i, j] = f'{dis_result[i, j]:.0f}\n\n{dis_var_result[i, j]:.0f}'  # 显示两个数值，每行一个    
        
        
        
            # PLOT
            fig,ax = plt.subplots(figsize=(15, 9),dpi=1500)
            ''' heatmap '''
            #ax = sns.heatmap(var[:15].T,annot=annot_str.T,
            #                 annot_kws={'size': 4},fmt="", cmap="YlGnBu", 
            #                 cbar_kws={'label': 'average RI location gap between obs and ensember members  Units: KM'}, 
            #                 linewidths=0.5,ax=ax)  #,vmin=0,vmax =30
            ax = sns.heatmap(var[:15].T,  # 删除了 annot=annot_str.T 参数
                 cmap="YlGnBu", 
                 cbar_kws={'label': 'average RI location gap between obs and ensember members  Units: KM'}, 
                 linewidths=0.5,ax=ax)  #,vmin=0,vmax =30
            ax.set_xlabel("Forecast Lead Time", fontsize=12)
            ax.set_ylabel("Forecast Initialization Time", fontsize=12)
            #  手动设置xticks数量,snsheatmap default上限是19个
            ax.set_yticks(np.arange(var.shape[1])) 
            ax.set_yticklabels(dates_name, rotation=0) 
            plt.xticks(rotation=0)
            xlabels=[str(i)+"d"  for i in range(1,16)]
            ax.set_xticklabels(xlabels)
            ax.set_title(ty+f"-RI{RIstd} average distance between RIobs and RImems   forecast_ini_mode:{ini_time_mode}" )
            cbar = ax.collections[0].colorbar
            cbar.locator = ticker.MaxNLocator(integer=True)
            cbar.update_ticks()
            
            
            ''' plot '''
            points = np.where(draw_data_obs[:,:15].T == 1)
            x_coords = points[0]+0.5
            y_coords = points[1]+0.15
            ax.scatter(x_coords, y_coords, color='red', marker='o',s=10)
            
            ax.invert_yaxis()
            if draw_opt==False:
                plt.savefig(os.path.join(pic_savepath, "{}.png".format(ty)), dpi=900)
                plt.close()
            else:
                plt.show()
                
            plt.clf()
            plt.close('all')
                
